# test_api.py
import requests
import json
import time
from datetime import datetime


class APITester:
    def __init__(self, base_url="http://localhost:8001"):
        self.base_url = base_url

    def test_root_endpoint(self):
        """测试根路径端点"""
        print("=== 测试根路径端点 ===")
        try:
            response = requests.get(self.base_url)
            if response.status_code == 200:
                data = response.json()
                print(f"✅ 成功: {data['message']}")
                print(f"版本: {data['version']}")
                print("可用端点:")
                for name, endpoint in data['endpoints'].items():
                    print(f"  - {name}: {endpoint}")
                return True
            else:
                print(f"❌ 失败: 状态码 {response.status_code}")
                return False
        except Exception as e:
            print(f"❌ 异常: {str(e)}")
            return False

    def test_health_endpoint(self):
        """测试健康检查端点"""
        print("\n=== 测试健康检查端点 ===")
        try:
            response = requests.get(f"{self.base_url}/health")
            if response.status_code == 200:
                data = response.json()
                print(f"✅ 健康状态: {data['status']}")
                print(f"时间戳: {data['timestamp']}")
                return True
            else:
                print(f"❌ 失败: 状态码 {response.status_code}")
                return False
        except Exception as e:
            print(f"❌ 异常: {str(e)}")
            return False

    def test_chat_endpoint(self):
        """测试聊天端点"""
        print("\n=== 测试聊天端点 ===")
        chat_data = {
            "message": "请介绍一下火星探测器",
            "history": [
                {
                    "type": "user",
                    "content": "你好",
                    "timestamp": datetime.now().isoformat()
                },
                {
                    "type": "assistant",
                    "content": "你好！我是火星科普助手，有什么我可以帮你的吗？",
                    "timestamp": datetime.now().isoformat()
                }
            ]
        }

        try:
            response = requests.post(
                f"{self.base_url}/api/chat2",
                json=chat_data,
                headers={"Content-Type": "application/json"}
            )

            if response.status_code == 200:
                data = response.json()
                print(f"✅ 聊天请求成功")
                print(f"成功状态: {data['success']}")

                # 显示完整回答内容
                print(f"完整回答:\n{data['message'][:500]}")

                # 显示相关问题
                if data['questions']:
                    print("\n相关问题:")
                    for i, question in enumerate(data['questions'], 1):
                        print(f"  {i}. {question}")
                else:
                    print("\n相关问题: 无")

                # 显示图片链接
                if data['images']:
                    print("\n相关图片:")
                    for i, image_url in enumerate(data['images'], 1):
                        print(f"  [{i}] {image_url}")
                else:
                    print("\n相关图片: 无")

                print(f"\n时间戳: {data['timestamp']}")
                return True
            else:
                print(f"❌ 失败: 状态码 {response.status_code}")
                print(f"错误详情: {response.text}")
                return False
        except Exception as e:
            print(f"❌ 异常: {str(e)}")
            return False

    def run_all_tests(self):
        """运行所有测试"""
        print("开始测试AI聊天助手API...")
        print("=" * 50)

        tests = [
            self.test_root_endpoint,
            self.test_health_endpoint,
            self.test_chat_endpoint
        ]

        results = []
        for test in tests:
            result = test()
            results.append(result)
            time.sleep(1)  # 避免请求过快

        print("\n" + "=" * 50)
        print("测试总结:")
        passed = sum(results)
        total = len(results)
        print(f"通过: {passed}/{total}")

        if passed == total:
            print("🎉 所有测试通过!")
        else:
            print("⚠️  部分测试失败，请检查服务状态")

        return passed == total


if __name__ == "__main__":
    # 创建测试实例
    tester = APITester()

    # 运行测试
    tester.run_all_tests()
